#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   data_table.py
@Time    :   2024/01/23 17:28:56
@Author  :   chakcy 
@Email   :   947105045@qq.com
@description   :   数据表
'''

import pandas as pd
import os
import sqlite3
import re


# 数据表
class DataTable():
    def __init__(self, file_path:str, sheet_name:str=None, sql=None):
        file_extension = os.path.splitext(file_path)[1]
        # csv 文件读取
        if file_extension == '.csv':
            self.data_table = pd.read_csv(file_path)
        # 读取 excel 文件
        elif (file_extension == '.xlsx') | (file_extension == '.xls'):
            self.data_table = pd.read_excel(file_path, sheet_name)
        # 读取 json 文件
        elif file_extension == '.json':
            self.data_table = pd.read_json(file_path)
        # 读取 sqlite 数据库
        elif file_extension == '.db':
            conn = sqlite3.connect(file_path)
            self.data_table = pd.read_sql(con=conn, sql=sql)
        else:
            raise ValueError("不支持的文件格式")

    # 获取数据
    def get_data_table(self):
        return self.data_table.to_dict("records")
    
# 筛选列
def filter_columns(list_dict:list, columns:list):
    data_table = pd.DataFrame(list_dict)
    return data_table[columns].to_dict("records")

# 筛选行
def filter_rows(list_dict:list, condition:str):
    """
    1. 筛选出某列的值等于某值的行：
        * filter_rows(list_dict, "x=='s'") # 筛选出x列为's'的行
    2. 多条件查询
        * filter_rows(list_dict, "x in ('s', 'a')") # 筛选出x列为's'或'a'的列
        * filter_rows(list_dict, "x in not ('s', 'a')") # 筛选出x列不为's'或'a'的列
    3. 比较数值列
        * filter_rows(list_dict, "x > 6 and y < 3") # 筛选出x列大于6并且y列大于3的行
    4. 列之间的比较
        * filter_rows(list_dict, "x < y") # 筛选出x列大于y列的行
    """
    data_table = pd.DataFrame(list_dict)
    return data_table.query(condition).to_dict("records")

# 模糊筛选
def fuzzy_filter(list_dict:list, column:str, pattern:str):
    data_table = pd.DataFrame(list_dict)
    regular = re.compile(pattern)
    return data_table.loc[data_table[column].str.contains(regular)].to_dict("records")

# 分组
def group_by(list_dict:list, column:str):
    data_table = pd.DataFrame(list_dict)
    return [[{"name": item[0], "data": item[1].to_dict("records")}] for item in data_table.groupby(by=column)]

# 增加列
def add_column(list_dict:list, column:str, data=None):
    data_table = pd.DataFrame(list_dict)
    if data is None:
        data_table.insert(len(data_table.columns), column, pd.NA)
    data_table.insert(len(data_table.columns), column, data)
    return data_table.to_dict("records")
    
if __name__ == '__main__':
    # dataTable = DataTable(file_path="test.xlsx", sheet_name="Sheet1")
    # dataTable = DataTable(file_path="test.csv")
    sql = """
    SELECT 
        *
    FROM
        `users`
    """
    dataTable = DataTable(file_path="example.db", sql=sql)
    # print(dataTable.data_table)
    data_table = dataTable.get_data_table()
    print(data_table)
    # print(filter_rows(data_table, "id==1 and name=='Alice'"))
    # print(add_column(data_table, "sex"))
    # print(fuzzy_filter(data_table, "姓名", "ch*"))
    print(group_by(data_table, "name"))
